
#include "reshape_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"// using namespace gert;/
#define BUFFER_NUM 2
#define BLOCK_SIZE 32

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  ReshapeTilingData tiling;
  // const gert::StorageShape* x1_shape = context->GetInputShape(0);
  // int32_t data_sz = 1;
  // for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
  //   data_sz *= x1_shape->GetStorageShape().GetDim(i);
  // tiling.set_size(data_sz);
  context->SetBlockDim(1);
  uint64_t ub_size;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  uint32_t totalLength = context->GetInputShape(0)->GetOriginShape().GetShapeSize();

  uint64_t typeSize = GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
  uint32_t preiter = ub_size/1/BUFFER_NUM/BLOCK_SIZE;
  int32_t blockElem = BLOCK_SIZE / typeSize;
  int32_t totalLengthAligned = (totalLength + blockElem - 1)/ blockElem * blockElem;
  uint32_t totalBlocks = totalLengthAligned / blockElem;
  // printf("totalLength:%d totalBlocks:%d preiter:%d blockElem:%d\n",totalLength,totalBlocks,preiter,blockElem);
  tiling.set_totalBlocks(totalBlocks);
  tiling.set_totalLength(totalLength);
  tiling.set_preiter(preiter);


  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {

// shape变换具体实现
template<typename T>
ge::graphStatus ReshapeInferShapeImpl(const T *reshape_dims, const gert::Shape &x_shape, gert::Shape &output_shape,int32_t xdim, int32_t reshape_rank, int32_t axes, int32_t num_axes) {
  constexpr T UNKNOWN_DIM = -1;
  // 将算子输出的维度数设置为reshape后的维度数reshape_rank
  // rank为确定的维度
  int32_t total =  xdim - num_axes + reshape_rank;
  output_shape.SetDimNum(total);
  int32_t x_shape_size = 1;
  // auto x_shape_size = x_shape.GetShapeSize();
  int64_t output_shapesize = 1;
  size_t unknown_dim_idx = std::numeric_limits<size_t>::max();
  for (int32_t i = 0; i < axes; i++) {
    output_shape.SetDim(i,x_shape.GetDim(i));
    output_shapesize *= x_shape.GetDim(i);
  }
  for (int32_t i = 0; i < reshape_rank; i++) {
    // printf("reshape_dims %d:%d\n",i,reshape_dims[i]);
    if (reshape_dims[i] != UNKNOWN_DIM) {  // reshape后某一轴的维度值不为-1 
      output_shape.SetDim(i+axes, reshape_dims[i]);  // 设置输出的维度值为reshape后的维度值
      output_shapesize *= reshape_dims[i];  // 计算当前输出元素数量
    } else {
      output_shape.SetDim(i+axes, 1);  // reshape后某一轴的维度值为-1，临时设置输出的维度值为1，后续计算后看是否可以推导出确定值，并记录未知维度的索引
      unknown_dim_idx = i+axes;
    }
  }
  for (int32_t i = axes+reshape_rank; i < total; i++) {
      output_shape.SetDim(i, reshape_dims[i-reshape_rank + num_axes]);
      output_shapesize *= x_shape.GetDim(i-reshape_rank + num_axes);
  }
  if (unknown_dim_idx == std::numeric_limits<size_t>::max() && output_shapesize == x_shape_size) {
    return ge::GRAPH_SUCCESS;  // 不存在未知维度，且输出shape size和输入x的shape size一致，直接返回成功
  } else if (unknown_dim_idx != std::numeric_limits<size_t>::max() && x_shape_size % output_shapesize == 0) {
    output_shape.SetDim(unknown_dim_idx, x_shape_size / output_shapesize); // 存在未知维度，根据输入shape动态调整未知维度值保持总元素个数不变
    return ge::GRAPH_SUCCESS;
  }
  return ge::GRAPH_SUCCESS;
}
// shape变换具体实现
// template<typename T>
static ge::graphStatus InferShape(gert::InferShapeContext* context) {
  const gert::Shape *x_shape = context->GetInputShape(0);        // 获取第0个输入的shape
  const gert::Tensor *shape_tensor = context->GetInputTensor(1); // 获取第1个输入的tensor
  const int64_t *axesVal = context->GetAttrs()->GetInt(0);
  int64_t axes = *axesVal;

  const int64_t *num_axesVal = context->GetAttrs()->GetInt(1);
  int64_t num_axes = *num_axesVal;

  int32_t xdim = x_shape->GetDimNum();

  gert::Shape *output_shape = context->GetOutputShape(0);
  if (x_shape == nullptr || shape_tensor == nullptr || output_shape == nullptr) {
      // 防御式编程，不应该出现的场景，打印错误并返回失败
      return ge::GRAPH_FAILED;
  }

  auto reshape_size = static_cast<int32_t>(shape_tensor->GetShapeSize());
  // if (num_axes != -1) {
  //     reshape_size = xdim - num_axes + reshape_size;
  // } else {
  //     num_axes = reshape_size;
  // }
  if (num_axes == -1) {
    num_axes = xdim;
  }
  if (reshape_size < 1) {
      // 防御式编程，不应该出现的场景，打印错误并返回失败
      return ge::GRAPH_FAILED;
  }

  // 根据原型信息，Reshape的shape输入支持INT32与INT64两类，根据不同的类型进入对应的模板函数中做真正的shape变换操作
  if (shape_tensor->GetDataType() == ge::DT_INT32) {
      const int32_t *reshape_data = shape_tensor->GetData<int32_t>();
      return ReshapeInferShapeImpl<int32_t>(reshape_data, *x_shape, *output_shape, xdim, reshape_size, axes, num_axes);
  } else {
      const int64_t *reshape_data = shape_tensor->GetData<int64_t>();
      return ReshapeInferShapeImpl<int64_t>(reshape_data, *x_shape, *output_shape, xdim, reshape_size, axes, num_axes);
  }
      return GRAPH_SUCCESS;

  }
    static graphStatus InferDataType(gert::InferDataTypeContext* context)
    {
        const auto inputDataType = context->GetInputDataType(0);
        context->SetOutputDataType(0, inputDataType);
        return ge::GRAPH_SUCCESS;
    }  
}


namespace ops {
class Reshape : public OpDef {
public:
    explicit Reshape(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("shape")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
          this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64, ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64, ge::DT_UINT8, ge::DT_UINT16, ge::DT_UINT32, ge::DT_UINT64})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("axis").Int();
        this->Attr("num_axes").Int();
        // this->SetInferShape(ge::InferShape);
        // this->SetInferDataType(ge::InferDataType); 
        // this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Reshape);
}
